
#include <ATen/ATen.h>
#include <ATen/Functions.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/library.h>

#include <THC/THC.h>
#include <THC/THCDeviceUtils.cuh>

#include <iostream>
#include <vector>


__global__ void add_H_W(float* bbox_pred_grad_mul_offset, const float* bbox_pred_grad_mul, const int* direction, const int* grad_mul_idx, const int C, const int CC, const int H, const int W) {
    const int row = blockIdx.x;
    const int col = blockIdx.y;
    const int yidx =  grad_mul_idx[col];

    for (int i=0; i<H; i++) {
        for (int j=0; j<W; j++) {
            int oidx = row * CC * H * W + col * H * W + i * W + j;
            int midx = row * C * H * W + yidx * H * W + i * W + j;
            bbox_pred_grad_mul_offset[oidx] = bbox_pred_grad_mul[midx] * direction[col];
        }
    }
}

at::Tensor star_dcn_offset_cuda(at::Tensor bbox_pred, 
                                at::Tensor dcn_base_offset, 
                                at::Tensor grad_mul_idx, 
                                at::Tensor direction,
                                int num_dconv_points, 
                                float gradient_mul, 
                                int stride) {

    AT_ASSERTM(bbox_pred.type().is_cuda(), "bbox_pred must be a CUDA tensor");
    const int N = bbox_pred.size(0);
    const int C = bbox_pred.size(1);
    const int H = bbox_pred.size(2);
    const int W = bbox_pred.size(3);

    // bbox_pred = bbox_pred.contiguous();
    // dcn_base_offset = dcn_base_offset.contiguous();
    // grad_mul_idx = grad_mul_idx.contiguous();
    // direction = direction.contiguous();
    int CC = 2 * num_dconv_points;

    // at::cuda::CUDAGuard device_guard(0);

    at::TensorOptions options =  dcn_base_offset.options();

    auto bbox_pred_grad_mul_offset = at::full({N, CC, H, W}, 0, options.dtype(at::kFloat));
    auto bbox_pred_grad_mul = ((1 - gradient_mul) * bbox_pred.detach() + gradient_mul * bbox_pred) / stride;


    float* bbox_pred_grad_mul_offset_ = bbox_pred_grad_mul_offset.data<float>();
    float* bbox_pred_grad_mul_ = bbox_pred_grad_mul.data<float>();
    int* direction_ = direction.data<int>();
    int* grad_mul_idx_ = grad_mul_idx.data<int>();

    dim3 grids(N, CC);
    dim3 blocks(1);
    add_H_W<<<grids, blocks>>>(
        bbox_pred_grad_mul_offset_,
        bbox_pred_grad_mul_,
        direction_,
        grad_mul_idx_,
        C,
        CC,
        H,
        W
    );
    
    return bbox_pred_grad_mul_offset - dcn_base_offset;
    // return dcn_base_offset;
}